[cuda backend] splitk turboquant sdpa for decode#20385
Merged
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20385
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 1 New Failure, 1 Unrelated FailureAs of commit 1af419b with merge base 39c0df6 ( NEW FAILURE - The following job has failed:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Gasoonjia
force-pushed
the
g4-128k-decode-splitk
branch
from
03a3725 to
c7273f7
Compare
mergennachin
approved these changes
Jun 22, 2026
Implements split-K (flash-decoding) decode path in tq4_sdpa to address severe under-parallelization at long context depths. The existing fused kernel launches only ~B*H_KV CTAs (e.g., 4 for gemma global GQA 8:4), each serially scanning all kv_len blocks. Changes: - Add _tq4_sdpa_decode_splitk_kernel: partitions KV across num_splits CTAs, each processing its chunk with TQ4 decompress + online softmax, writing partial O/M/L to scratch tensors - Add _tq4_sdpa_decode_reduce_kernel: reduces partials across splits with proper rescaling by exp(m_split - m_global) - Add _launch_tq4_decode_splitk: host launcher with static num_splits for CUDA-graph compatibility (derived from buffer size, not runtime kv_len) - Modify tq4_sdpa(): dispatch to split-K when L_q==1 AND kv_len>=256, preserving existing behavior for prefill and short decode - Add unit tests: test_kv_len_clamp_decode_*_splitk verifies split-K output matches both fp32 reference (cosine ~0.99+) AND fused kernel output, with garbage tail negative control for kv_len bound correctness CUDA-graph constraints satisfied: - num_splits static at capture time (from Lk buffer size) - Scratch tensors pre-allocated with fixed addresses - kv_len read on-device via tl.load (no host sync) - No dynamic allocation in captured region Test results: - test_kv_len_clamp_decode_gemma_global_splitk: PASSED (D=512, GQA 8:4, kv_len=8192/32768 with 32k buffer and garbage tail) - test_kv_len_clamp_decode_qwen_splitk: PASSED (D=256, GQA 16:2, kv_len=8192/32768 with 32k buffer and garbage tail) - All existing tq4_sdpa tests remain green
The split-K decode path was branching on runtime kv_len value via kv_len_t.item() >= threshold, which is data-dependent control flow that breaks torch.export/AOTI tracing. Fix: Make dispatch static by using N_KV (buffer size) instead of runtime kv_len value. The kv_len tensor is still used inside the kernel for bounds checking via tl.load on device (CUDA-graph safe). - Dispatch now based on static N_KV >= 256 (buffer size) - Removes incorrect guard_or_false usage with .item() - Preserves kv_len on-device usage in kernel for bounds - Maintains byte-identical behavior for qwen (kv_len fallback) - Eager correctness preserved
…ot fixed 64) The split-K decode launcher hardcoded BLOCK_M=64 for the pack-GQA path, but only L_q*num_groups rows are valid (=8 for Gemma4 global, H_Q=32/H_KV=4). That left 56/64 M-rows idle while the QK/PV tensor-core MMAs still computed all 64, and made acc[64, HEAD_DIM=512] fp32 = 128 KB/CTA, spilling registers to local memory. Size BLOCK_M = max(16, next_pow2(L_q*num_groups)) instead (16 for Gemma4 global), removing the wasted MMA rows and the spills. Decode-only; prefill path unchanged.
- gemma4_31b cuda_source_transformations: pass mask_is_causal=True to tq4_sdpa (the 12th arg was missing). Enables the prefill per-tile causal block-skip (~halves causal-triangle work); no-op for decode. - tq4_sdpa kernel + tests: drop model-specific (Gemma/Qwen) references; describe configs by shape (head_dim / GQA ratio) instead. - tests: add split-K coverage for B=2 (test_splitk_batch2) and a non-contiguous query stride (test_splitk_noncontiguous_query).
Gasoonjia
force-pushed
the
g4-128k-decode-splitk
branch
from
c7273f7 to
1af419b
Compare
This PR needs a
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Gemma4-31b decode performance dramastically drop when decode length is large, blocking it to be a local agent engine.
This PR introduces the split-k version turboquant sdpa for decode to introduce more parallelism and better utilize SM cores and other GPU hardware resource to accelerate turboquant-supported long context mode like Gemma4-31b.
With splitk-sdpa, we can mantain the perf in the short number of decode tokens, while dramastically improve the perf by 2.77x for deoce length == 32768, and perf is still steady even if decode length == 120k
test env: A100, export seq length == 128k, turboquant enabled